﻿using System;
using System.Net.Sockets;

namespace DotNetClient
{
    public enum TransportState
    {
        readHead = 1,   // on read head
        readBody = 2,   // on read body
        closed = 3      // connection closed, will ignore all the message and wait for clean up
    }

    class StateObject
    {
        public const int BufferSize = 1024;
        internal byte[] buffer = new byte[BufferSize];
    }

    public class Transporter
    {
        public const int HEADER_LENGTH = 5;

        private Socket _socket;
        private Action<Package> _messageProcesser;

        // Used for get message
        private StateObject _stateObject = new StateObject();
        private TransportState _transportState;
        private IAsyncResult _asyncReceive;
        private IAsyncResult _asyncSend;
        private bool _onSending = false;
        private bool _onReceiving = false;

        private byte[] _headBuffer = new byte[HEADER_LENGTH];
        private byte[] _msgBuffer;
        private int _bufferOffset = 0;
        private int _msgLength = 0;

        internal Action onDisconnect = null;

        private object _lock = new object();

        public Transporter(Socket socket, Action<Package> processer)
        {
            _socket = socket;
            _messageProcesser = processer;
            _transportState = TransportState.readHead;
        }

        public void start()
        {
            receive();
        }

        internal void close()
        {
            _transportState = TransportState.closed;
            try
            {
                if (_onReceiving) _socket.EndReceive(_asyncReceive);
                if (_onSending) _socket.EndSend(_asyncSend);
            }
            catch (Exception e)
            {
                UnityEngine.Debug.LogError(e.Message);
            }
        }

        public void send(byte[] buffer)
        {
            if (_transportState != TransportState.closed)
            {
                _asyncSend = _socket.BeginSend(buffer, 0, buffer.Length, SocketFlags.None, new AsyncCallback(sendCallback), _socket);
                _onSending = true;
            }
        }

        private void sendCallback(IAsyncResult asyncSend)
        {
            if (_transportState == TransportState.closed) return;

            _socket.EndSend(asyncSend);
            _onSending = false;
        }

        public void receive()
        {
            _asyncReceive = _socket.BeginReceive(_stateObject.buffer, 0, _stateObject.buffer.Length, SocketFlags.None, new AsyncCallback(endReceive), _stateObject);
            _onReceiving = true;
        }

        private void endReceive(IAsyncResult asyncReceive)
        {
            if (_transportState == TransportState.closed) return;

            StateObject state = (StateObject)asyncReceive.AsyncState;
            try
            {
                int length = _socket.EndReceive(asyncReceive);
                //UnityEngine.Debug.Log("EndReceive: " + length);
                _onReceiving = false;
                if (length > 0)
                {
                    processBytes(state.buffer, 0, length);
                    // Receive next message
                    if (_transportState != TransportState.closed)
                    {
                        receive();
                    }
                }
                else
                {
                    onDisconnect?.Invoke();
                }
            }
            catch (SocketException)
            {
                onDisconnect?.Invoke();
            }
        }

        internal void processBytes(byte[] bytes, int offset, int limit)
        {
            if (_transportState == TransportState.readHead)
            {
                readHead(bytes, offset, limit);
            }
            else if (_transportState == TransportState.readBody)
            {
                readBody(bytes, offset, limit);
            }
        }

        /// <summary>
        /// Head
        /// </summary>
        /// <param name="bytes"></param>
        /// <param name="offset"></param>
        /// <param name="limit"></param>
        /// <returns></returns>
        private bool readHead(byte[] bytes, int offset, int limit)
        {
            int length = limit - offset;
            int headNum = HEADER_LENGTH - _bufferOffset;

            if (length >= headNum)
            {
                // Write head buffers
                writeBytes(bytes, 0, HEADER_LENGTH, _headBuffer);
                // Get package length
                _msgLength = ((_headBuffer[0] << 24) | (_headBuffer[1] << 16) | (_headBuffer[2] << 8) | _headBuffer[3]) - 1;

                // Init message buffer
                _msgBuffer = new byte[_msgLength];
                offset += headNum;
                _bufferOffset = 0;
                _transportState = TransportState.readBody;

                if (offset <= limit) processBytes(bytes, offset, limit);
                return true;
            }
            else
            {
                writeBytes(bytes, offset, length, _bufferOffset, _headBuffer);
                _bufferOffset += length;
                return false;
            }
        }

        /// <summary>
        /// Body
        /// </summary>
        /// <param name="bytes"></param>
        /// <param name="offset"></param>
        /// <param name="limit"></param>
        private void readBody(byte[] bytes, int offset, int limit)
        {
            if ((offset + _msgLength) <= limit)
            {
                // Get package type
                PackageType packageType = (PackageType)bytes[HEADER_LENGTH - 1];
                // Write msg buffers
                writeBytes(bytes, offset, _msgLength, _bufferOffset, _msgBuffer);
                offset += _msgLength;

                // Invoke the protocol api to handle the message
                Package package = new Package(packageType, _msgBuffer);
                _messageProcesser.Invoke(package);
                _bufferOffset = 0;
                _msgLength = 0;

                if (_transportState != TransportState.closed)
                    _transportState = TransportState.readHead;
                if (offset < limit)
                    processBytes(bytes, offset, limit);
            }
            else
            {
                writeBytes(bytes, offset, limit - offset, _bufferOffset, _msgBuffer);
                _bufferOffset += limit - offset;
                _transportState = TransportState.readBody;
            }
        }

        private void writeBytes(byte[] source, int start, int length, byte[] target)
        {
            writeBytes(source, start, length, 0, target);
        }

        private void writeBytes(byte[] source, int start, int length, int offset, byte[] target)
        {
            for (int i = 0; i < length; i++)
            {
                target[offset + i] = source[start + i];
            }
        }
    }
}
